#pragma warning(disable : 4996)
#include <iostream>
#include <list>
#include <cstdio>
#include <cmath>
#include <stack>
#include <queue>
#include <cstring>
#include <sstream>
#include <set>
#include <string>
#include <iterator>
#include <vector>
#include <map>
#include <algorithm>

#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image.h"
#include "stb_image_write.h"

using namespace std;

// ����char����8λһȡ�������������ַ����ͻ��Ǻ��ֻ���������ʽ�����Խ���ѹ��
const unsigned int N = 256;   // char����8λ����ôһ�������256�����ܵ��ַ���Ҷ�ڵ���Ŀ
const unsigned int M = 2 * N; // ����ӵ��n��Ҷ����huffman����һ����2*n-1����㣬�������һλ����������select����ѡ��ʹ��
const unsigned long long MAX = 0xffffffffffff;

class Buffer
{ // �����࣬���ļ����ж�ȡ/д�������ʱ��ͨ��Buffer����������ַ��Ķ�ȡ/д��
public:
    unsigned int bits; // ʵ���ֽڳ���
    char ch;           // �ֽ�
};

class HuffmanTree; // huffman��������

class treeNode
{ // huffman�������
    friend HuffmanTree;

private:
    unsigned long long weight; // �ý��һ�������˶��ٴ�
    unsigned int right, left;
    unsigned int parents;
};

class HuffmanTree
{ // huffman����
private:
    treeNode nodes[M];     // �洢����ÿ�������Ϣ�����M��
    unsigned int leaf[N];  // �洢Ҷ�ڵ��ַ���Ϣ�����N��
    unsigned int index[N]; // �洢Ҷ�ڵ���±�ֵ
    char *leafCode[N];     // �洢Ҷ�ڵ�ı�����Ϣ
    FILE *input, *output;
    unsigned int num;        // ʵ����Ҷ������
    unsigned long long size; // һ���ж��ٸ��ַ�
    Buffer buf;              // �����������

    // ��������
    void write(unsigned int i);                      // ���ļ���д��һ��bite
    void write(unsigned int num, unsigned int bits); // ���ļ���д��bitsλ��num��
    void writerest();                                // ������buf���滹���ַ�û��д�룬���������һ��д��
    void read(unsigned int &i);                      // ���ļ��ж�ȡһλbite����i���
    void read(unsigned int &num, unsigned int bits); // ���ļ��ж�ȡbitsλ����������num���

    // Encode��������
    void enSieve();                                                    // ��ȡ�ļ�������ͳ���ļ����ַ���
    void select(unsigned int pos, unsigned int &t1, unsigned int &t2); // ��0����pos��ѡ������Ȩ����С�ģ���t1��t2���
    void enSetTree();                                                  // ����Ҫ��ѹ�����ļ��е��ַ�����huffman������������ÿ�����ı�����Ϣ

    // Decode��������
    void deSieve(); // ��ȡ��Ҫ����ѹ�ļ������������������

public:
    void Encode(); // ���ļ�����
    void Decode(); // ���ļ�����
};

// ����������
void HuffmanTree::write(unsigned int i)
{                               // ���ļ���д��һ��bite
    buf.ch = (buf.ch << 1) + i; // ��ch����һλ
    buf.bits++;                 // ch�İ�λbite��ʵ��ʹ������һ
    if (buf.bits == 8)
    { // ���ȫ�����꣬��ô��buf.ch�����output�У�����buf.bits,buf.ch��0���¿�ʼ��һ�ּ���
        fputc(buf.ch, output);
        buf.bits = 0;
        buf.ch = 0;
    }
}
void HuffmanTree::write(unsigned int target, unsigned int bits)
{ // ���ļ���д��bitsλ��num��
    /*ɵ������
    for (unsigned int i = 0; i < bits; ++i) {//һλһλ�ķŽ�ȥ
        write((target & 128)>>7 );//�Ӹ�λ���ηŵ���λ
        target <<= 1;
    }
    */

    stack<unsigned int> s;
    unsigned int i, bit;
    for (i = 1; i <= bits; i++)
    {
        s.push(target & 1);
        target = (target >> 1);
    }
    for (i = 1; i <= bits; i++)
    {
        bit = s.top();
        write(bit);
        s.pop();
    }
}
void HuffmanTree::writerest()
{ // ������buf���滹���ַ�û��д�룬���������һ��д��
    unsigned int now = buf.bits;
    if (now > 0)
        for (unsigned int i = 0; i < 8 - now; i++)
            write(0);
}
void HuffmanTree::read(unsigned int &i)
{ // ���ļ��ж�ȡһλbite����i���
    if (buf.bits == 0)
    {
        buf.bits = 8;
        buf.ch = fgetc(input);
    }
    i = (buf.ch & 128) >> 7;
    buf.bits--;
    buf.ch <<= 1;
}
void HuffmanTree::read(unsigned int &target, unsigned int bits)
{ // ���ļ��ж�ȡbitsλ����������num���
    unsigned int tmp;
    target = 0;
    for (unsigned int i = 0; i < bits; ++i)
    {
        read(tmp);
        target = (target << 1) + tmp;
    }
}

// Encode����������
void HuffmanTree::enSieve()
{ // ��ȡ�ļ�������ͳ���ļ����ַ���
    char inName[1000], outName[1000];
    cout << "Input file name that you want to code:";
    cin >> inName;
    cout << "Input target file name:";
    cin >> outName;
    if ((input = fopen(inName, "rb")) == NULL)
    {
        cout << "Can not open file." << endl;
        system("pause");
        exit(1);
    }
    if (feof(input))
    {
        cout << "Empty source file" << endl;
        system("pause");
        exit(1);
    }
    if ((output = fopen(outName, "wb")) == NULL)
    {
        cout << "Can not open file." << endl;
        ;
        system("pause");
        exit(1);
    }

    // ���ļ��ж�ȡ�ַ�����ͳ���ַ�����Ƶ��
    rewind(input);
    unsigned int ch;
    size = 0;
    for (unsigned int i = 0; i < N; ++i)
    {
        leaf[i] = 0;
        index[i] = 0;
    }
    for (unsigned int i = 0; i < M; ++i)
    {
        nodes[i].weight = 0;
        nodes[i].left = nodes[i].right = nodes[i].parents = M - 1;
    }
    ch = fgetc(input);
    while (!feof(input))
    {
        leaf[ch]++;
        size++;
        ch = fgetc(input);
    }

    // nodes[N-1].weight��Ϊ���
    nodes[M - 1].weight = MAX;

    // ɸ������Ƶ��Ϊ0���ַ�����д��nodes��index���飬���޸�numֵ
    num = 0;
    for (unsigned int i = 0; i < N; ++i)
        if (leaf[i])
        {
            nodes[num].weight = leaf[i];
            leaf[i] = num;
            index[num] = i;
            num++;
        }
    if (!num)
    {
        cout << "doesn't have a word" << endl;
        system("pause");
        exit(1);
    }
}

void HuffmanTree::select(unsigned int pos, unsigned int &t1, unsigned int &t2)
{ // ��0����pos��ѡ������Ȩ����С�ģ���t1��t2���
    t1 = M - 1, t2 = M - 1;
    for (unsigned int i = 0; i < pos; ++i)
    {
        if (nodes[i].weight < nodes[t1].weight && nodes[i].parents == M - 1)
            t1 = i;
    }
    for (unsigned int i = 0; i < pos; ++i)
    {
        if (nodes[i].weight < nodes[t2].weight && i != t1 && nodes[i].parents == M - 1)
            t2 = i;
    }
}

void HuffmanTree::enSetTree()
{ // ����Ҫ��ѹ�����ļ��е��ַ�����huffman������������ÿ�����ı�����Ϣ
    // ����huffman��
    for (unsigned int i = num; i < num * 2 - 1; ++i)
    {
        unsigned int t1, t2;
        select(i, t1, t2);
        nodes[i].weight = nodes[t1].weight + nodes[t2].weight;
        nodes[i].left = t1;
        nodes[i].right = t2;
        nodes[t1].parents = nodes[t2].parents = i;
    }

    for (int i = 0; i < N; ++i)
        leafCode[i] = NULL;

    // ��ÿ�������б���
    unsigned int start, c, f, i;
    char *cd = new char[num]; // ������ʱ����
    for (i = 0; i < N; i++)
        if (leafCode[i] != NULL)
        {
            delete[] leafCode[i]; // �ͷŴ洢�ռ�
            leafCode[i] = NULL;
        }
    cd[num - 1] = '\0'; // ���������
    for (i = 0; i < num; i++)
    {                    // ��λ��Huffman����
        start = num - 1; // ���������λ��
        for (c = i, f = nodes[i].parents; f != M - 1; c = f, f = nodes[c].parents)
        { // ��Ҷ���������
            if (nodes[f].left == c)
                cd[--start] = '0';
            else
                cd[--start] = '1';
        }
        leafCode[i] = new char[num - start]; // Ϊ��i���ַ��������ռ�
        strcpy(leafCode[i], &cd[start]);     // ��cd���Ʊ��뵽HuffmanCode
    }
    delete cd;
}

void HuffmanTree::Encode()
{                // ���ļ�����
    enSieve();   // ��ʼ��input��output��ͳ���ļ����ַ�
    enSetTree(); // ����enSieve���huffman���Ľ�������ַ����б���

    rewind(output);
    rewind(input);
    // ��output�Ŀ�ͷ��д�����ṹ
    buf.bits = 0;
    buf.ch = 0;
    fwrite(&size, sizeof(unsigned long long), 1, output); // д��size
    write(num, 8);                                        // �����ṹ�е�Ҷ������д��
    for (unsigned int i = 0; i < num; ++i)                // �����ڵ��е�Ҷ�ڵ�д��
        fwrite(&index[i], sizeof(char), 1, output);
    // ѡ��num�����Ҫ����λ���洢
    unsigned maxbit = 1;
    unsigned int tmp = num * 2 - 1;
    while (tmp)
    {
        maxbit++;
        tmp >>= 1;
    }
    for (unsigned int i = num; i < num * 2 - 1; ++i)
    { // д�����Һ�����Ϣ
        write(nodes[i].left, maxbit);
        write(nodes[i].right, maxbit);
    }

    // д�������Ϣ
    unsigned int ch;
    ch = fgetc(input);
    while (!feof(input))
    {
        unsigned int start = 0; // �ж϶�ch�ı���leafCode[loc]����ʼλ��
        while (leafCode[leaf[ch]][start] != '\0')
        {
            if (leafCode[leaf[ch]][start] == '1')
                write(1);
            else
                write(0);
            ++start;
        }
        ch = fgetc(input);
    }
    writerest(); // д��ʣ�µ��ַ�
    cout << "Done!\n\n";
    fclose(input);
    fclose(output);
}

// Decode����������
void HuffmanTree::deSieve()
{ // ��ȡ��Ҫ����ѹ�ļ������������������
    char inName[1000], outName[1000];
    cout << "Input file name that you want to decode:";
    cin >> inName;
    cout << "Input target file name:";
    cin >> outName;
    if ((input = fopen(inName, "rb")) == NULL)
    {
        cout << "Can not open file." << endl;
        system("pause");
        exit(1);
    }
    if (feof(input))
    {
        cout << "Empty source file" << endl;
        system("pause");
        exit(1);
    }
    if ((output = fopen(outName, "wb")) == NULL)
    {
        cout << "Can not open file." << endl;
        ;
        system("pause");
        exit(1);
    }

    // ��ʼ��ȡ���ṹ
    rewind(input);
    for (unsigned int i = 0; i < M; ++i)
    {
        nodes[i].parents = nodes[i].right = nodes[i].left = N - 1;
    }

    buf.bits = 0; // ��ջ�����
    buf.ch = 0;
    fread(&size, sizeof(unsigned long long), 1, input); // ��ȡsize
    read(num, 8);                                       // ��ȡ���ṹ�е�Ҷ������
    if (num == 0)
        num = 256;
    for (unsigned int i = 0; i < num; ++i) // ��ȡ���ڵ��е�Ҷ�ڵ�
        fread(&index[i], sizeof(char), 1, input);
    // ѡ��num�����Ҫ����λ���洢
    unsigned maxbit = 1;
    unsigned int tmp = num * 2 - 1;

    while (tmp)
    {
        maxbit++;
        tmp >>= 1;
    }

    for (unsigned int i = num; i < num * 2 - 1; ++i)
    { // ��ȡ���Һ�����Ϣ
        read(nodes[i].left, maxbit);
        read(nodes[i].right, maxbit);
        nodes[nodes[i].left].parents = nodes[nodes[i].right].parents = i;
    }
}

void HuffmanTree::Decode()
{
    deSieve();

    // ��ʼ����
    rewind(output);
    unsigned int tmp;
    read(tmp);
    for (int i = 0; i < size; ++i)
    {
        unsigned int loc = 2 * num - 2;
        while ((nodes[loc].left != N - 1 || nodes[loc].right != N - 1) && !feof(input))
        {
            if (tmp == 0)
                loc = nodes[loc].left;
            else
                loc = nodes[loc].right;
            read(tmp);
        }
        fputc(index[loc], output);
    }
    cout << "Done!\n\n";
    fclose(input);
    fclose(output);
}

class ImageCompress
{
public:
    // 压缩图片,quality范围1-100,1为最低质量,100为最高质量
    static bool compress(const string &input_name, const string &output_name, int quality)
    {
        int width, height, channels;
        unsigned char *image_data = stbi_load(input_name.c_str(), &width, &height, &channels, 0);
        if (!image_data)
        {
            return false;
        }

        bool success = false;
        string ext = get_file_extension(output_name);

        if (ext == "jpg" || ext == "jpeg")
        {
            success = stbi_write_jpg(output_name.c_str(), width, height, channels, image_data, quality);
        }
        else if (ext == "png")
        {
            // 对于PNG，我们先将质量转换为压缩级别（1-100）
            // 然后根据质量等级调整图像数据
            unsigned char *compressed_data = new unsigned char[width * height * channels];

            // 根据quality调整图像质量
            float compression_ratio = (100 - quality) / 100.0f;
            for (int i = 0; i < width * height * channels; i++)
            {
                // 对每个像素进行量化，质量越低，量化步长越大
                int step = 1 + (int)(compression_ratio * 7); // 1-8的量化步长
                compressed_data[i] = (image_data[i] / step) * step;
            }

            // 使用压缩后的数据写入PNG
            success = stbi_write_png(output_name.c_str(), width, height, channels, compressed_data, width * channels);
            delete[] compressed_data;
        }
        else if (ext == "bmp")
        {
            success = stbi_write_bmp(output_name.c_str(), width, height, channels, image_data);
        }
        else if (ext == "tga")
        {
            success = stbi_write_tga(output_name.c_str(), width, height, channels, image_data);
        }

        stbi_image_free(image_data);
        return success;
    }

private:
    static string get_file_extension(const string &path)
    {
        size_t dot_pos = path.find_last_of(".");
        if (dot_pos != string::npos)
        {
            string ext = path.substr(dot_pos + 1);
            // 转换为小写
            for (char &c : ext)
            {
                c = tolower(c);
            }
            return ext;
        }
        return "";
    }
};

int main(void)
{
    char choose = '1';
    while (choose != '4')
    {
        cout << "1.Huffman Encode" << endl;
        cout << "2.Huffman Decode" << endl;
        cout << "3.Image Compress" << endl;
        cout << "4.exit" << endl;
        cin >> choose;

        switch (choose)
        {
        case '1':
        {
            HuffmanTree tree;
            tree.Encode();
            break;
        }
        case '2':
        {
            HuffmanTree tree;
            tree.Decode();
            break;
        }
        case '3':
        {
            string input_name, output_name;
            int quality;
            cout << "Input image name (with extension): ";
            cin >> input_name;
            cout << "Output image name (with extension): ";
            cin >> output_name;
            cout << "Compression quality (1-100): ";
            cin >> quality;

            if (quality < 1)
                quality = 1;
            if (quality > 100)
                quality = 100;

            if (ImageCompress::compress(input_name, output_name, quality))
            {
                cout << "Image compressed successfully!" << endl;
            }
            else
            {
                cout << "Failed to compress image." << endl;
            }
            break;
        }
        default:
            break;
        }
    }
    return 0;
}